from sklearnex import patch_sklearn

patch_sklearn()
import numpy as np
from skimage.feature import hog, local_binary_pattern, sift
from sklearn.svm import SVC
from sklearn import metrics
import mnist_reader


# 提取hog特征
def extract_hog_features(data):
    hog_features = []
    count = 0
    for img in data:
        count = count + 1
        print(f"\rExtracting HOG features: {count}/{len(data)}", end='')
        hog_feature = hog(img, pixels_per_cell=(8, 8), cells_per_block=(2, 2))
        hog_features.append(hog_feature)
    print("\n")
    return np.array(hog_features)


# 提取LBP特征
def extract_lbp_features(data, P=8, R=1):
    lbp_features = []
    count = 0
    for img in data:
        count = count + 1
        print(f"\rExtracting LBP features: {count}/{len(data)}", end='')
        lbp = local_binary_pattern(img, P, R, method='default')
        # We use the histogram of LBP as the feature vector
        (hist, _) = np.histogram(lbp.ravel(), bins=np.arange(0, P + 3), range=(0, P + 2))
        hist = hist.astype("float")
        hist /= (hist.sum() + 1e-6)  # Normalize the histogram
        lbp_features.append(hist)
    print("\n")
    return np.array(lbp_features)


def extract_sift_features(data):
    sift_features = []
    count = 0
    for img in data:
        count += 1
        print(f"\rExtracting SIFT features: {count}/{len(data)}", end='')
        keypoints, descriptors = sift.FeatureDetector()
        if descriptors is None:
            descriptors = np.zeros((1, 128))  # Use zero descriptors if no keypoints are found
        sift_features.append(descriptors.flatten())
    print("\n")
    return np.array(sift_features)


# 提取像素特征
def extract_pixel_features(data):
    pixel_features = []
    count = 0
    for img in data:
        count = count + 1
        print(f"\rExtracting Pixel features: {count}/{len(data)}", end='')
        # 将图像展平成一维数组，并将像素值归一化到[0, 1]
        pixel_feature = img.flatten() / 255.0
        pixel_features.append(pixel_feature)
    print("\n")
    return np.array(pixel_features)


# 结果评估
def evaluate(y_val, y_pred):
    # 结果评估
    matrix = metrics.confusion_matrix(y_val, y_pred)  # 混淆矩阵
    print("confusion_metrix:\n", matrix)

    accuracy = metrics.accuracy_score(y_val, y_pred)  # accuracy精度
    precision = metrics.precision_score(y_val, y_pred, average='macro')  # precision
    recall = metrics.recall_score(y_val, y_pred, average='macro')  # recall
    f1 = metrics.f1_score(y_val, y_pred, average='macro')  # f1 score
    print("Accuracy:", accuracy, " Precision:", precision, " Recall:", recall, " F1:", f1)

    # # 绘制ROC曲线 牛魔的多分类画不了
    # fpr, tpr, thresholds = metrics.roc_curve(y_val, y_pred)
    # plt.figure(figsize=(10, 6))
    # plt.xlim(0, 1)  # 设定x轴的范围
    # plt.ylim(0.0, 1.1)  # 设定y轴的范围
    # plt.xlabel('False Postive Rate')
    # plt.ylabel('True Postive Rate')
    # plt.plot(fpr, tpr, linewidth=2, linestyle="-", color='red')
    # plt.show()


if __name__ == '__main__':
    # 定义svm模型
    '''
        训练可以选择的核函数有:
        ’linear’：线性核函数
        ‘poly’：多项式核函数
        ‘rbf’：径像核函数/高斯核
        ‘sigmoid’：sigmoid核函数
        还有一个‘precomputed’：核矩阵 不知道怎么用
        '''
    svm_model = SVC(kernel='sigmoid', verbose=True)
    # 1.读取数据集
    x_train, y_train = mnist_reader.load_mnist('data', kind='train')
    x_test, y_test = mnist_reader.load_mnist('data', kind='t10k')
    print(len(x_train), len(y_train), len(x_test), len(y_test))

    # ==========使用hog进行训练测试==========
    # 2.调整图片大小形状
    x_train_hog = x_train.reshape(-1, 28, 28)
    x_test_hog = x_test.reshape(-1, 28, 28)

    # 3.提取特征
    x_train_hog_features = extract_hog_features(x_train_hog)
    x_test_hog_features = extract_hog_features(x_test_hog)

    # 4.训练模型
    print("hog start training...")
    svm_model.fit(x_train_hog_features, y_train)
    # 预测
    print("hog start predicting...")
    y_pred_hog = svm_model.predict(x_test_hog_features)

    # 结果评估
    evaluate(y_test, y_pred_hog)

    # ==========使用lbp进行训练==========
    # 2.调整图片大小形状
    x_train_lbp = x_train.reshape(-1, 28, 28)
    x_test_lbp = x_test.reshape(-1, 28, 28)

    # 3.提取特征
    x_train_lbp_features = extract_lbp_features(x_train_lbp)
    x_test_lbp_features = extract_lbp_features(x_test_lbp)

    # 4.训练
    print("lbp start training...")
    svm_model.fit(x_train_lbp_features, y_train)
    # 5.预测
    print("lbp start predicting...")
    y_pred_lbp = svm_model.predict(x_test_lbp_features)

    # 6.结果评估
    evaluate(y_test, y_pred_lbp)

    # ==========使像素特征进行训练测试==========
    x_train_pixel = x_train.reshape(-1, 28, 28)
    x_test_pixel = x_test.reshape(-1, 28, 28)

    # 3.提取特征
    x_train_pixel_features = extract_pixel_features(x_train_pixel)
    x_test_pixel_features = extract_pixel_features(x_test_pixel)

    # 4.训练
    print("pixel start training...")
    svm_model.fit(x_train_pixel_features, y_train)
    # 5.预测
    print("pixel start predicting...")
    y_pred_pixel = svm_model.predict(x_test_pixel_features)

    # 6.结果评估
    evaluate(y_test, y_pred_pixel)
